import random
import time
import numpy as np
import pygame
import math
from enum import Enum
from typing import Dict, Tuple, Sequence
import multiprocessing as mp
import ctypes
import zmq
import struct
import socket
import pickle
from Env import *
# from Env_PF import *
from utils.utils import *

######## processing 1 - Recv Ctrl Order By ZMQ ########
def Ctrl_Receiver(_recv_dict : dict, recv_port : int):
    context = zmq.Context(1)
    recv_socket = context.socket(zmq.PULL)
    recv_socket.connect("tcp://127.0.0.1:" + str(recv_port))

    while True:
        time.sleep(0.001)
        try:
            recv_pkg = recv_socket.recv(zmq.NOBLOCK)
            recv_msg = pickle.loads(recv_pkg)
            for i in range(len(recv_msg)):
                _recv_dict[recv_msg[i][0]] = recv_msg[i][1:]
            # print("Recv Action!", _recv_dict)
        except:
            continue
        


##### processing 2 - Send to Server Block ############
def Data_Sender(_send_dict : dict, send_port : int, _update : mp.Value):  # TODO: 根据Srv的状态进行修改
    # Initialize a zeromq context
    context = zmq.Context(1)
    # Set up a channel to send work
    send_socket = context.socket(zmq.PUSH)
    send_socket.bind("tcp://127.0.0.1:" + str(send_port))
 
    while True:
        send_msg = []
        if _update.value == 0:
            # time.sleep(0.3)
            continue
        
        for key, value in _send_dict.items():
            send_msg.append(value)
        try:
            send_pkg = pickle.dumps(send_msg)
            send_socket.send(send_pkg, zmq.NOBLOCK)
            # print("Send From MVS", send_msg)
            _update.value = 0   # 【重要！！！】保证数据同步
        except:
            # print("No Receiver!")
            continue

########### processing 3 - Simulation Scenario ############
def Scenario(_recv_dict : dict, _send_dict : dict, init_msgs : InitMsgs, _update : mp.Value):
    # Init Vehicle State(type, seq, pos, ang, spd)
    # 画布初始化
    screen_init = pygame.display.set_mode(init_msgs.screen_init_msg['sc_size'])
    screen_init.fill(init_msgs.screen_init_msg['bg_color'])
    pygame.display.set_caption("General Observation Map")
    
    # 坐标转换参数初始化
    coord_convert = init_msgs.screen_init_msg['coord_convert']
    
    # Init Vehicle Message
    if len(init_msgs.vehicle_init_msg['params']) == 0:
        print("No Vehicle Init Message! Failed to Init Environment!")
        return -1
    vehicles, targets = {}, {}
    # TODO：计算各机spawn_pos
    for i, msg in enumerate(init_msgs.vehicle_init_msg['params']):
        # 实例化 vehicle 类
        spawn_pos = convert_coordinates((msg[1], msg[2]),(coord_convert[0],coord_convert[1]), is_flip=True) # dx, dy,
        Veh = Vehicle(default_init=msg, default_spawn=(spawn_pos[0,0], spawn_pos[0,1]), leader_mode=True if msg[0]==1 else False)  # change to your vehicle c
        vehicles[msg[0]] = Veh
        # 实例化 target 类  # formation offset
        
        Tar = Target()
        targets[msg[0]] = Tar
        
    general_scenario = GeneralObservationMap(screen_init=screen_init, 
                                             vehicles=vehicles,
                                             formation_set=init_msgs.decision_cfg['formation_cfg'])
    lock = _update.get_lock()
    cur_eps = 0
    count = 0
    while True:
        
        for event in pygame.event.get():
            # print(1)
            if event.type == pygame.QUIT:
                general_scenario.close()
        
        general_scenario.render()
        lock.acquire()
        try:
            _update.value = 0
            # action, state, reward, done 均为 dict 参数表，在环境类中批量更新
            # a = general_scenario.sample_actions()   # TODO: recv_msg
            a = _recv_dict
            s, r, done, terminate = general_scenario.step(action=a)   # TODO: send_msg
            
            for i, params in enumerate(init_msgs.vehicle_init_msg['params']):
                # print(s[params[0]])
                # print(r[params[0]])
                # print(done[params[0]])
                # print(terminate[params[0]])
                _send_dict[params[0]] = [params[0]] + s[params[0]] + [r[params[0]]] + [done[params[0]]] + [terminate[params[0]]] + [count] + [cur_eps] # id, state, next_state, reward, terminate, steps(步长)[steps=0则reset], cur_eps
            
            if all(val for val in terminate.values()) or a == {}:
                s = general_scenario.reset()
                print(count)
                cur_eps += 1
                print(cur_eps)
                count = 0
                
            count +=1        
            # print(_send_dict)
            _update.value = 1
        finally:
            lock.release()    
       
        
                            
########### main ############
if __name__ == "__main__":
    # 读取配置文件
    init_msgs = process_yaml_cfg(path_to_configs='/root/code/workspace/src/MVS/Configs/', filename='InitConfigs.yaml')
    recv_port = init_msgs.render_ports_cfg['recv_port']
    send_port = init_msgs.render_ports_cfg['send_port']
    
    # for port in [recv_port, send_port]:
    #     procs = find_procs_by_port(port)
    #     if procs:
    #         for proc in procs:
    #             print(f"Found process {proc.pid} occupying port {port}. Attempting to kill.")
    #             try:
    #                 kill_proc_tree(proc.pid)
    #             except psutil.NoSuchProcess:
    #                 pass
    #     else:
    #         print(f"No process found occupying port {port}.")

    # 数据收发类型
    mgr = mp.Manager()
    _recv_dict = mgr.dict()
    _send_dict = mgr.dict()
    
    # 数据更新标志位
    _update = mp.Value('i', 0)
    
    pc1 = mp.Process(target=Ctrl_Receiver, args=(_recv_dict, recv_port, ))
    pc1.start()
    
    pc2 = mp.Process(target=Data_Sender, args=(_send_dict, send_port, _update, ))
    pc2.start()
    
    pc3 = mp.Process(target=Scenario, args=(_recv_dict, _send_dict, init_msgs, _update, ))
    pc3.start()

    pc1.join()
    pc2.join()
    pc3.join()